算法 树状数组|线段树

树状数组 O(logn)

  1. 在某个位置上的数加上一个数, 单点修改
  2. 求任意某区间的前缀和, 区间查询

如果想用前缀和实现单点修改+区间查询, 需要O(n) + O(1) = O(n), 而树状数组中两种操作都是O(logn).

  • tr[x] = ( x - lowbit(x), x ] , 每个tr[x]代表了这个一个区间内的前缀和

  • 三大操作 : lowbit + add + query

  • query(x) 查询的是x前的所有前缀和

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
// 单点修改 + 区间查询
#include<iostream>
#define rep(i, a, b) for(int i = (a); i <= (b); i ++ )
using namespace std;
const int N = 100010;
int n, m;
int a[N], tr[N];

int lowbit(int x)
{
    return x & -x;
}

void add(int x, int v)
{
    for(int i = x; i <= n; i += lowbit(i)) tr[i] += v;
}

int query(int x)
{
    int ret = 0;
    for(int i = x; i > 0; i -= lowbit(i)) ret += tr[i];
    return ret;
}

signed main()
{
    cin >> n >> m;
    rep(i, 1, n) cin >> a[i];
    rep(i, 1, n) add(i, a[i]);
    while(m -- )
    {
        int op, x, y; cin >> op >> x >> y;
        if(op == 0) cout << query(y) - query(x - 1) << endl;
        else add(x, y);
    }
    return 0;
}

区间修改 + 单点查询

原来树状数组维护的是原数组a, query得到的结果的前缀和s, 整体过程是 a -> s. 那么 b->a 也有相同的积分关系, 也就是差分数组到原数组, 我们可以将add改为两个点的修改操作(b), query得到的结果是a, 这就是单点的值.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
// 区间修改 + 单点查询
#include<iostream>
#define rep(i, a, b) for(int i = (a); i <= (b); i ++)
using namespace std;
const int N = 100010;
int n, m;
int a[N];
int tr[N * 2];

int lowbit(int x)
{
    return x & -x;
}

void add(int x, int v)
{
    for(int i = x; i < N; i += lowbit(i)) tr[i] += v;
}

int query(int x)
{
    int res = 0;
    for(int i = x; i > 0; i -= lowbit(i)) res += tr[i];
    return res;
}

int main()
{
    cin >> n >> m;
    rep(i, 1, n) cin >> a[i];
    rep(i, 1, n) add(i, a[i] - a[i - 1]); // tr要维护的是差分数组
    
    while(m -- )
    {
        string op; cin >> op;
        if(op == "C")
        {
            int l, r, d;
            cin >> l >> r >> d;
            add(l, d), add(r + 1, -d);  // 用两个单点操作取代区间操作
        }
        else 
        {
            int x;
            cin >> x;
            cout << query(x) << endl;	// 单点查询求a
        }
    }
    return 0;
}

区间修改 + 区间查询

这种操作实际希望我们达到的效果是 b -> s, 也就是我们通过修改差分数组, 之后就可以通过O(logn)的时间得到前缀和, 想要实现b->s, 是通过分析数学公式得到的 :

image-20250323101825181

这里的d可以当作b, 我们看图可以分析出 si = (d1 + d2 + … + dn) * (x + 1) - (1 * d1 + 2 * d2 + .. + n * dn);

前半部和后半部要求我们各维护一个树状数组, 整体过程可以理解为b —树状数组O(logn)—> 特殊的a —数学计算O(1)—> s.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
// 区间修改 + 区间查询
#include<iostream>
#define rep(i, a, b) for(int i = (a); i <= (b); i ++ )
#define int long long 
using namespace std;
const int N = 100010;
int n, m;
int a[N];
int tr1[N * 2], tr2[N * 2];

int lowbit(int x)
{
    return x & -x;
}

void add(int tr[], int x, int c) // 修改b
{
    for(int i = x; i <= N; i += lowbit(i)) tr[i] += c;
}

int query(int tr[], int x) // 获取a
{
    int res = 0;
    for(int i = x; i > 0; i -= lowbit(i)) res += tr[i];
    return res;
}

int sum(int x)	// 通过a再获取s
{
    return query(tr1, x) * (x + 1) - query(tr2, x);
}

signed main()
{
    cin >> n >> m;
    rep(i, 1, n) cin >> a[i];
    rep(i, 1, n) 
    {
        int b = a[i] - a[i - 1];
        add(tr1, i, b);
        add(tr2, i, i * b);
    }
    
    while(m -- )
    {
        string op;
        int l, r;
        cin >> op >> l >> r;
        if(op == "C")
        {
            int d; cin >> d;
            add(tr1, l, d), add(tr2, l, l * d);
            add(tr1, r + 1, -d), add(tr2, r + 1, (r + 1) * -d);
        }
        else 
        {
            cout << sum(r) - sum(l - 1) << endl;
        }
    }
    return 0;
}

树状数组 + 二分答案

这种算法一般是 遍历 + 树状数组 + 二分, 时间复杂度是O(nlognlogn), 其实非常快捷, 1e5 - 1e6的数据都可以解决.

该算法常用来获取第k大元素, 相比于小根堆, 其优势在于k值非常灵活, 小根堆的k值必须固定, 但是这种算法的k可以是任意的.

<L3-002 特殊堆栈> 堆栈是一种经典的后进先出的线性结构，相关的操作主要有“入栈”（在堆栈顶插入一个元素）和“出栈”（将栈顶元素返回并从堆栈中删除）。本题要求你实现另一个附加的操作：“取中值”——即返回所有堆栈中元素键值的中值。给定 N 个元素，如果 N 是偶数，则中值定义为第 N/2 小元；若是奇数，则为第 (N+1)/2 小元。

该题在普通栈的基础上可以求中值, 中值在这里近似于第k大元素, k 可以通过计算得来, 是灵活可变的. 并且这里入堆的值只到1e5, 于是我们就可以用树状数组维护[1 - 1e5]区间上在堆中的元素个数, push就单点修改 + 1, pop就单点修改 - 1, query(x)得到的就是堆中有的范围在[1, x] 中的元素个数.

有了这样的树状数组, 当我需要求第k大元素时, 我们可以发现, 假如我们二分答案, 只有在mid大于等于第k大元素时query(mid)的结果才会大于等于k, 这个过程是单增的可以二分.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
#include<iostream>
#include<vector>
#include<cstring>
#include<algorithm>
#include<stack>
#define rep(i, a, b) for(int i = (a); i <= (b); i ++ )
using namespace std;
const int N = 100010;
// 求第k大元素(k可灵活) -> 树状数组 + 二分
int n;
int tr[N * 2];
stack<int> st;

int lowbit(int x)
{
    return x & -x;
}

void add(int x, int v)
{
    for(int i = x; i <= N; i += lowbit(i)) tr[i] += v;
}

int query(int x)
{
    int res = 0;
    for(int i = x; i > 0; i -= lowbit(i)) res += tr[i];
    return res;
}

signed main()
{
    cin >> n;
    rep(i, 1, n)
    {
        string op;
        cin >> op;
        if(op == "Pop")
        {
            if(!st.size()) cout << "Invalid" << endl;
            else 
            {
                cout << st.top() << endl;
                add(st.top(), -1);
                st.pop();
            }
        }
        else if(op == "PeekMedian")
        {
            if(!st.size()) cout << "Invalid" << endl;
            else 
            {
                // 计算中值k
                int sz = st.size();
                int k;
                if(sz & 1) k = (sz + 1) / 2;
                else k = sz / 2;
                // 二分答案, 找第k个元素的元素值
                int l = 1, r = N;
                while(l < r)
                {
                    int mid = l + r >> 1;
                    if(query(mid) >= k) r = mid;
                    else l = mid + 1;
                }
                cout << l << endl;
            }
        }
        else
        {
            int x; cin >> x;
            st.push(x);
            add(x, 1);
        }
    }
    return 0;
}

有 n 头奶牛，已知它们的身高为 1∼n 且各不相同，但不知道每头奶牛的具体身高。

现在这 n 头奶牛站成一列，已知第 i 头牛前面有 Ai 头牛比它低，求每头奶牛的身高。

这道题就更意识流了, 没有告诉你每头奶牛对应的值, 只是告诉你范围在[1, n]和每头奶牛左边比其低的个数.

首先要从边界考虑, 我从右往左思考, 第n头奶牛右边包含所有的牛, 那么a[n] + 1一定代表其名次, 并且由于范围是[1, n]并且不重样, 那么第n头奶牛的值就是a[n] + 1.

我们继续考虑n - 1, 其名次其实是受n影响的, 也就是说a[n - 1] + 1不一定可以代表n - 1的名次的, 但是有一种方法可以让我们确定n - 1的名次, 那就是树状数组 + 二分, 这里该算法做到的是 :

  • 树状数组维护区间[1, n], 表示区间上未被使用的值有多少个, 从右往左遍历, 和上题一样二分答案, 找到第a[i] + 1大的值, 这个值就是第i头的值 , 一旦确定了第i头的值, 就把该值从tr中删去, 这样在以后的遍历中, i右边的牛就不会再影响到结果了.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
// 树状数组 + 二分
#include<iostream>
#define rep(i, a, b) for(int i = (a); i <= (b); i ++ )
#define int long long 
using namespace std;
const int N = 100010;
int n;
int a[N];
int tr[N * 2];
int ans[N];

int lowbit(int x)
{
    return x & -x;
}

void add(int x, int c)
{
    for(int i = x; i < N; i += lowbit(i)) tr[i] += c;
}

int query(int x)
{
    int res = 0;
    for(int i = x; i > 0; i -= lowbit(i)) res += tr[i];
    return res;
}

signed main()
{
    cin >> n;
    rep(i, 2, n) cin >> a[i];
    
    rep(i, 1, n) add(i, 1);
    
    for(int i = n; i > 0; i -- )
    {
        int rk = a[i] + 1;
        int l = 1, r = n;
        while(l < r) 
        {
            int mid = l + r >> 1;
            if(query(mid) >= rk) r = mid;
            else l = mid + 1;
        }
        add(l, -1);
        ans[i] = l;
    }
    
    rep(i, 1, n) cout << ans[i] << endl;
    return 0;
}

线段树

作用依旧是单点修改 + 区间查询, 但是相比于树状数组只能对数值进行维护, 线段树可以维护各种区间的性质, 比如最大值, 最小值, 最长连续子序列和等等.

  • pushup : 用子节点信息更新当前节点
  • build : 在一段区间上初始化线段树
  • modify : 修改
  • query : 查询
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
#include<iostream>
#define rep(i, a, b) for(int i = (a); i <= (b); i ++ )
using namespace std;
const int N = 100010;
int n, m;
int a[N];
struct node
{
    int l, r;
    int sum;
}tr[N * 4];

void pushup(int u)
{
    tr[u].sum = tr[u << 1].sum + tr[u << 1 | 1].sum;
}

void build(int u, int l, int r)
{
    if(l == r)
    {
        tr[u] = {l, r, a[l]};
        return;
    }
    tr[u] = {l, r};
    int mid = tr[u].l + tr[u].r >> 1;
    build(u << 1, tr[u].l, mid);
    build(u << 1 | 1, mid + 1, tr[u].r);
    pushup(u);
}

int query(int u, int l, int r)
{
    if(tr[u].l >= l && tr[u].r <= r) return tr[u].sum;
    int res = 0;
    int mid = tr[u].l + tr[u].r >> 1;
    if(l <= mid) res += query(u << 1, l, r);
    if(r > mid) res += query(u << 1 | 1, l, r);
    return res;
}

void modify(int u, int x, int v)
{
    if(tr[u].l == tr[u].r && tr[u].l == x) 
    {
        tr[u].sum += v;
        return;
    }
    int mid = tr[u].l + tr[u].r >> 1;
    if(x <= mid) modify(u << 1, x, v);
    else modify(u << 1 | 1, x, v);
    pushup(u);
}

signed main()
{
    cin >> n >> m;
    rep(i, 1, n) cin >> a[i];
    build(1, 1, n);
    
    while(m -- )
    {
        int op, x, y;
        cin >> op >> x >> y;
        if(op == 0) cout << query(1, x, y) << endl;
        else modify(1, x, y);
    }
}
#算法
算法 树状数组|线段树
http://example.com/2025/03/29/[算法] 树状数组线段树/
作者
天目中云
发布于
2025年3月29日
许可协议